import os
import json
from pathlib import Path
from typing import Dict, List

import cv2
import google.generativeai as genai
import numpy as np
import PIL.Image

CONFIG: Dict = {
    "model_name": "models/gemini-2.0-flash-exp",
    "defect_classes": ["Corrosion", "Crack", "Peeled Paint", "Damage", "Coating Degradation"],
    "use_proxy": True,
    "proxy_host": "127.0.0.1",
    "proxy_port": "7890",
    "enable_preprocess": True,
    "min_confidence": 0.65,
    "min_box_area": 0.0001,
    "max_box_area": 0.92,
    "merge_nearby_defects": True,
    "merge_distance_threshold": 0.10,
}

if CONFIG["use_proxy"]:
    os.environ.setdefault("HTTP_PROXY", f"http://{CONFIG['proxy_host']}:{CONFIG['proxy_port']}")
    os.environ.setdefault("HTTPS_PROXY", f"http://{CONFIG['proxy_host']}:{CONFIG['proxy_port']}")


def preprocess_image(image: PIL.Image.Image) -> PIL.Image.Image:
    """增强图像以突出缺陷特征，特别是锈蚀和涂层剥落"""
    if not CONFIG["enable_preprocess"]:
        return image.convert("RGB")

    rgb = np.array(image.convert("RGB"))
    
    # 双边滤波：保留边缘的同时去噪
    denoised = cv2.bilateralFilter(rgb, d=5, sigmaColor=30, sigmaSpace=30)
    
    # LAB色彩空间增强：突出颜色差异
    lab = cv2.cvtColor(denoised, cv2.COLOR_RGB2LAB)
    l, a, b = cv2.split(lab)
    
    # L通道（亮度）：增强涂层剥落的亮区对比
    clahe_l = cv2.createCLAHE(clipLimit=3.8, tileGridSize=(8, 8))
    l = clahe_l.apply(l)
    
    # A通道（绿-红）：增强锈蚀的红色/褐色特征
    clahe_a = cv2.createCLAHE(clipLimit=3.5, tileGridSize=(6, 6))
    a = clahe_a.apply(a)
    
    # B通道（蓝-黄）：增强锈蚀的黄褐色特征
    clahe_b = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(6, 6))
    b = clahe_b.apply(b)
    
    # 合并增强后的通道
    enhanced = cv2.cvtColor(cv2.merge((l, a, b)), cv2.COLOR_LAB2RGB)
    
    # 锐化：突出纹理和边界
    blur = cv2.GaussianBlur(enhanced, (0, 0), sigmaX=1.2)
    sharpened = cv2.addWeighted(enhanced, 1.8, blur, -0.8, 0)
    
    # 轻微饱和度提升：使锈色更明显
    hsv = cv2.cvtColor(sharpened, cv2.COLOR_RGB2HSV)
    h, s, v = cv2.split(hsv)
    s = cv2.multiply(s, 1.15)
    s = np.clip(s, 0, 255).astype(np.uint8)
    final = cv2.cvtColor(cv2.merge((h, s, v)), cv2.COLOR_HSV2RGB)
    
    return PIL.Image.fromarray(final)


def build_prompt() -> str:
    return f"""You are an expert wind turbine blade inspector with 20+ years of experience in composite material defect analysis. Your mission is to create precise training annotations for a YOLO object detection model used in critical safety inspections.

=== BACKGROUND CONTEXT ===
Wind turbine blades are composite structures (fiberglass/carbon fiber) with protective gel coat. They operate in harsh environments (UV, rain, hail, lightning, temperature cycles) causing various surface defects. Early detection prevents catastrophic failures.

=== SYSTEMATIC INSPECTION PROTOCOL ===

STEP 1: INITIAL SCAN
- Mentally divide the blade into a 5x5 grid
- Scan each grid cell systematically (left→right, top→bottom)
- Note the baseline "clean blade" appearance (uniform white/light gray)

STEP 2: ANOMALY DETECTION
For each grid cell, ask:
- Does color differ from baseline? (brown/rust/dark patches OR bright white exposed areas)
- Does texture differ? (rough/granular OR smooth exposed substrate)
- Are there linear discontinuities? (cracks/gaps)
- Are there 3D deformations? (dents/gouges)

STEP 3: DEFECT CLASSIFICATION

**CORROSION (锈蚀)**
Visual Signature:
  - Color: Brownish, rust-orange, reddish-brown, tan, or dark brown staining
  - Texture: Granular, rough, pitted, or corroded appearance
  - Pattern: Irregular blotches, often radiating from metal components
  - Context: Frequently near fasteners, leading edge strips, or exposed metal inserts

Critical Distinction:
  - Has warm color tones (brown/orange/rust)
  - Surface appears degraded/rough
  - Color penetrates into material (not surface coating)

Real-world example: Rust bleeding from hidden fastener, brownish water stain pattern

**PEELED PAINT (涂层剥落)**
Visual Signature:
  - Color: Exposed underlayer appears LIGHTER/WHITER/BRIGHTER than intact coating
  - Texture: Sharp, irregular, jagged flaking edges
  - Pattern: Islands of missing coating with clear boundaries
  - Context: High stress areas, impact zones, edges

Critical Distinction:
  - Has cool/neutral tones (white/gray/light)
  - Shows substrate exposure (different layer visible)
  - Has distinct edge where coating flaked off

Real-world example: Paint chip revealing white composite underneath

**CRACK (裂纹)**
Visual Signature:
  - Appearance: Thin dark lines showing material separation
  - Geometry: Linear, branching (spider-web), or curved
  - Depth: Visible shadow or depth indication
  - Pattern: Continuous lines (may be hairline to wide)

Critical Distinction:
  - NOT straight panel seams (those are manufacturing joints)
  - Shows actual material discontinuity
  - Has irregular path (not perfectly straight)

Real-world example: Lightning strike crack, fatigue crack from cyclic loading

**DAMAGE (损伤)**
Visual Signature:
  - Form: Gouges, dents, impact craters, missing chunks
  - Texture: Exposed fibers (stringy white/gray material visible)
  - Depth: 3D deformation or material loss evident
  - Context: Impact points, collision areas

Critical Distinction:
  - Shows physical material removal or deformation
  - Often has irregular shapes
  - May show underlying fiber structure

Real-world example: Hail impact dent, bird strike damage

**COATING DEGRADATION (涂层退化)**
Visual Signature:
  - Appearance: Chalky, powdery, or matte surface finish
  - Pattern: Large uniform areas (not localized spots)
  - Texture: Surface roughness without material loss
  - Effect: Loss of gloss, micro-cracking networks

Critical Distinction:
  - Widespread gradual deterioration (not sudden damage)
  - Surface layer only (no substrate exposure)
  - Uniform rather than patchy

Real-world example: UV aging, general weathering over large area

=== CRITICAL EXCLUSION CHECKLIST ===
DO NOT annotate these normal features:
✗ Panel seams: Straight lines with regular spacing (manufacturing joints)
✗ Bolts/Fasteners: Circular metal components (intentional hardware)
✗ Lightning protection strips: Metallic tape strips (intentional feature)
✗ Drainage holes: Small circular openings (design feature)
✗ Shadows: Dark areas from 3D features (lighting artifact, no texture change)
✗ Dirt/Water stains: Removable surface contamination (wipes clean)
✗ Glare/Reflections: Bright spots from lighting (no actual defect)
✗ Text/Logos: Intentional markings

=== DECISION TREE ===
For each suspicious region:
1. Is it a normal feature? (see exclusion list) → SKIP
2. Does it have brown/rust color? → Likely CORROSION
3. Does it show bright exposed substrate? → Likely PEELED PAINT
4. Is it a linear gap with depth? → Likely CRACK
5. Is there material loss/deformation? → Likely DAMAGE
6. Is it widespread surface wear? → Likely COATING DEGRADATION
7. Still uncertain? → SKIP (be conservative)

=== BOUNDING BOX PRECISION ===
- Pixel-tight: Box edges within 2-3 pixels of defect boundary
- Minimal background: Every pixel inside box should contribute to defect identification
- Normalized coordinates: All values between 0.0 and 1.0
- Validity: Ensure x1 < x2 and y1 < y2

=== CONFIDENCE SCORING ===
- 0.90-1.00: Crystal clear defect, perfect lighting, no ambiguity
- 0.80-0.89: Clear defect, good visibility, minor boundary uncertainty
- 0.70-0.79: Visible defect, adequate lighting, moderate uncertainty
- 0.60-0.69: Subtle defect, challenging conditions (shadows/distance/blur)
- Below 0.60: DO NOT ANNOTATE (too uncertain)

=== OUTPUT REQUIREMENTS ===
Return ONLY a valid JSON array. NO explanations. NO markdown. NO additional text.

Format:
[
  {{"defect_type": "Corrosion", "bounding_box": {{"x1": 0.123, "y1": 0.234, "x2": 0.156, "y2": 0.267}}, "confidence": 0.87}},
  {{"defect_type": "Peeled Paint", "bounding_box": {{"x1": 0.345, "y1": 0.456, "x2": 0.378, "y2": 0.489}}, "confidence": 0.92}}
]

If no defects found: []

=== FINAL INSTRUCTIONS ===
1. Apply the systematic protocol above
2. Be thorough but conservative
3. Prioritize precision over recall
4. Return ONLY valid JSON

Begin inspection now.
"""


def call_model(image: PIL.Image.Image) -> List[Dict] | None:
    """调用AI模型进行缺陷检测，优化参数以提高准确性"""
    try:
        model = genai.GenerativeModel(
            model_name=CONFIG["model_name"],
            generation_config={
                "temperature": 0.1,  # 低温度：更确定性的输出
                "top_p": 0.85,  # 核采样：专注于高概率tokens
                "top_k": 32,  # 限制候选tokens：减少随机性
                "max_output_tokens": 8192,  # 支持大量检测结果
            },
        )
        
        response = model.generate_content([build_prompt(), image])
        text = (response.text or "").strip()
        
        # 清理JSON格式
        text = text.replace("```json", "").replace("```", "").strip()
        
        # 提取JSON数组
        if text and not text.startswith("["):
            start_idx = text.find("[")
            text = text[start_idx:] if start_idx != -1 else ""
        
        if text and not text.endswith("]"):
            end_idx = text.rfind("]")
            text = text[:end_idx + 1] if end_idx != -1 else ""
        
        if not text:
            return []
        
        # 解析JSON
        result = json.loads(text)
        
        # 验证返回的是列表
        if not isinstance(result, list):
            print(f"警告: API返回的不是列表格式")
            return []
        
        return result
        
    except json.JSONDecodeError as e:
        print(f"JSON解析失败: {e}")
        print(f"原始响应前200字符: {text[:200] if 'text' in locals() else '无'}")
        return None
    except Exception as exc:
        print(f"模型调用失败: {exc}")
        return None


def filter_defects(items: List[Dict]) -> List[Dict]:
    results = []
    for item in items or []:
        if not isinstance(item, dict) or item.get("defect_type") not in CONFIG["defect_classes"]:
            continue
        if "bounding_box" not in item or "confidence" not in item:
            continue
        
        try:
            conf = float(item["confidence"])
            if conf < CONFIG["min_confidence"]:
                continue
            
            box = item["bounding_box"]
            x1, y1, x2, y2 = float(box["x1"]), float(box["y1"]), float(box["x2"]), float(box["y2"])
            
            if not (0.0 <= x1 < x2 <= 1.0 and 0.0 <= y1 < y2 <= 1.0):
                continue
            
            area = (x2 - x1) * (y2 - y1)
            if not (CONFIG["min_box_area"] <= area <= CONFIG["max_box_area"]):
                continue
            
            results.append({
                "defect_type": item["defect_type"],
                "bounding_box": {"x1": x1, "y1": y1, "x2": x2, "y2": y2},
                "confidence": conf,
            })
        except (KeyError, TypeError, ValueError):
            continue
    
    return results


def merge_nearby_defects(defects: List[Dict]) -> List[Dict]:
    if not CONFIG["merge_nearby_defects"] or not defects:
        return defects
    
    by_type = {}
    for d in defects:
        by_type.setdefault(d["defect_type"], []).append(d)
    
    merged = []
    threshold = CONFIG["merge_distance_threshold"]
    
    for dtype, dlist in by_type.items():
        n = len(dlist)
        parent = list(range(n))
        
        def find(x):
            if parent[x] != x:
                parent[x] = find(parent[x])
            return parent[x]
        
        def union(x, y):
            px, py = find(x), find(y)
            if px != py:
                parent[px] = py
        
        for i in range(n):
            for j in range(i + 1, n):
                b1, b2 = dlist[i]["bounding_box"], dlist[j]["bounding_box"]
                cx1, cy1 = (b1["x1"] + b1["x2"]) / 2, (b1["y1"] + b1["y2"]) / 2
                cx2, cy2 = (b2["x1"] + b2["x2"]) / 2, (b2["y1"] + b2["y2"]) / 2
                if ((cx1 - cx2)**2 + (cy1 - cy2)**2)**0.5 < threshold:
                    union(i, j)
        
        groups = {}
        for i in range(n):
            groups.setdefault(find(i), []).append(dlist[i])
        
        for group in groups.values():
            if len(group) == 1:
                merged.append(group[0])
            else:
                merged.append({
                    "defect_type": dtype,
                    "bounding_box": {
                        "x1": min(d["bounding_box"]["x1"] for d in group),
                        "y1": min(d["bounding_box"]["y1"] for d in group),
                        "x2": max(d["bounding_box"]["x2"] for d in group),
                        "y2": max(d["bounding_box"]["y2"] for d in group),
                    },
                    "confidence": sum(d["confidence"] for d in group) / len(group),
                })
    
    return merged


def save_yolo(defects: List[Dict], image_path: Path) -> None:
    lines = []
    for d in defects:
        idx = CONFIG["defect_classes"].index(d["defect_type"])
        b = d["bounding_box"]
        xc, yc = (b["x1"] + b["x2"]) / 2, (b["y1"] + b["y2"]) / 2
        w, h = b["x2"] - b["x1"], b["y2"] - b["y1"]
        lines.append(f"{idx} {xc:.6f} {yc:.6f} {w:.6f} {h:.6f}")
    image_path.with_suffix(".txt").write_text("\n".join(lines), encoding="utf-8")


def draw_boxes(defects: List[Dict], image_path: Path, output_path: Path) -> None:
    img = cv2.imread(str(image_path))
    if img is None:
        return
    
    h, w = img.shape[:2]
    colors = {"Corrosion": (0, 0, 255), "Crack": (255, 0, 0), "Peeled Paint": (0, 255, 255),
              "Damage": (0, 165, 255), "Coating Degradation": (0, 255, 0)}
    
    for d in defects:
        b = d["bounding_box"]
        x1, y1 = int(b["x1"] * w), int(b["y1"] * h)
        x2, y2 = int(b["x2"] * w), int(b["y2"] * h)
        color = colors.get(d["defect_type"], (255, 255, 255))
        
        cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
        
        label = f"{d['defect_type']} {d['confidence']:.2f}"
        (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
        cv2.rectangle(img, (x1, y1 - th - 6), (x1 + tw + 4, y1), color, -1)
        cv2.putText(img, label, (x1 + 2, y1 - 3), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
    
    cv2.imwrite(str(output_path), img)


def main() -> None:
    if not (api_key := os.environ.get("GOOGLE_API_KEY")):
        print("❌ 未检测到环境变量 GOOGLE_API_KEY")
        return
    
    genai.configure(api_key=api_key)
    
    image_path = Path("DJI_0013_03_04_png.rf.b4c7d90d73dabcc44badf5baef29b7eb.jpg")
    if not image_path.exists():
        print(f"❌ 未找到图片: {image_path}")
        return
    
    print("=" * 60)
    print("风力发电叶片缺陷检测系统")
    print("=" * 60)
    
    print("\n[1/5] 正在预处理图像...")
    print("      - 双边滤波去噪")
    print("      - LAB色彩空间增强（突出锈蚀和剥落）")
    print("      - 自适应锐化")
    print("      - 饱和度优化")
    processed = preprocess_image(PIL.Image.open(image_path))
    print("      ✓ 图像预处理完成")
    
    print("\n[2/5] 正在调用AI模型进行智能分析...")
    print("      - 系统性网格扫描")
    print("      - 多层决策树分类")
    print("      - 精确边界框定位")
    raw = call_model(processed)
    
    if raw is None:
        print("      ❌ AI模型调用失败")
        return

    print(f"      ✓ AI返回 {len(raw)} 个候选区域")
    
    if raw:
        print("\n[3/5] 原始检测结果（前10项）:")
        for i, item in enumerate(raw[:10], 1):
            dtype = item.get('defect_type', '?')
            conf = item.get('confidence', 0)
            bbox = item.get('bounding_box', {})
            area = (bbox.get('x2', 0) - bbox.get('x1', 0)) * (bbox.get('y2', 0) - bbox.get('y1', 0))
            print(f"      {i:2d}. {dtype:20s} | 置信度:{conf:.2f} | 面积:{area:.4f}")
        if len(raw) > 10:
            print(f"      ... 还有 {len(raw) - 10} 个候选区域")
    
    print("\n[4/5] 正在过滤和合并...")
    defects = filter_defects(raw)
    print(f"      - 置信度过滤: {len(defects)} 个通过")
    
    if not defects:
        print("      ⚠ 未检测到符合标准的缺陷")
        image_path.with_suffix(".txt").write_text("", encoding="utf-8")
        return

    if CONFIG["merge_nearby_defects"]:
        before_merge = len(defects)
        defects = merge_nearby_defects(defects)
        print(f"      - 邻近合并: {before_merge} → {len(defects)} 个")
    
    print("\n[5/5] 生成标注文件...")
    save_yolo(defects, image_path)
    draw_boxes(defects, image_path, image_path.with_name(f"{image_path.stem}_annotated.jpg"))
    
    print("\n" + "=" * 60)
    print("检测完成！最终统计:")
    print("=" * 60)
    
    counts = {}
    for d in defects:
        counts[d["defect_type"]] = counts.get(d["defect_type"], 0) + 1
    
    total = sum(counts.values())
    for dtype in CONFIG["defect_classes"]:
        count = counts.get(dtype, 0)
        if count > 0:
            bar = "█" * int(count / total * 30)
            print(f"  {dtype:20s} : {count:2d} 个 {bar}")
    
    print(f"\n  总计: {total} 个缺陷")
    print(f"\n  输出文件:")
    print(f"    - YOLO标注: {image_path.with_suffix('.txt').name}")
    print(f"    - 可视化图: {image_path.stem}_annotated.jpg")
    print("=" * 60)


if __name__ == "__main__":
    main()
